using System.Collections.Generic;

namespace Ast.Core.Services.Extensions
{
    public static class PsLoadedModuleListServiceExtensions
    {
        public static KernelModule GetNtModule(this IPsLoadedModuleListService psLoadedModuleListService)
        {
            KernelModule retVal = null;
            psLoadedModuleListService.ListModules((entry, path, name) =>
            {
                retVal = new KernelModule
                {
                    Name = name,
                    Path = path,
                    StartAddress = (uint) entry.DllBase,
                    EndAddress = (uint) entry.DllBase + entry.SizeOfImage,
                    EntryAddress = (uint) entry.EntryPoint
                };
                return false;
            });
            return retVal;
        }

        public static IList<KernelModule> GetAllModules(this IPsLoadedModuleListService psLoadedModuleListService)
        {
            var lst = new List<KernelModule>();
            psLoadedModuleListService.ListModules((entry, path, name) =>
            {
                lst.Add(new KernelModule
                {
                    Name = name,
                    Path = path,
                    StartAddress = (uint) entry.DllBase,
                    EndAddress = (uint) entry.DllBase + entry.SizeOfImage,
                    EntryAddress = (uint) entry.EntryPoint
                });
                return true;
            });
            return lst;
        }
    }

    public class KernelModule
    {
        public string Name { get; set; }
        public string Path { get; set; }
        public uint StartAddress { get; set; }
        public uint EndAddress { get; set; }
        public uint EntryAddress { get; set; }
    }
}